from typing import List, Dict
import heapq
import random
import time
import pygame
from models import Position, Vehicle, Direction, TrafficLightStatus
from visualization import TrafficVisualization
# 导入参数系统
from parameters import get_params

# 获取参数
PARAMS = get_params()

class Road:
    def __init__(self, start: Position, end: Position):
        self.start = start
        self.end = end
        self.length = ((end.x - start.x) ** 2 + (end.y - start.y) ** 2) ** 0.5
        self.vehicles: List[Vehicle] = []
        # 使用参数系统中的道路容量
        self.capacity = PARAMS["道路"]["capacity"]  # 道路容量
        # 添加速度限制
        self.speed_limit = PARAMS["道路"]["default_speed_limit"]

    @property
    def density(self) -> float:
        return len(self.vehicles) / self.capacity


class Intersection:
    def __init__(self, position: Position):
        self.position = position
        self.traffic_light = TrafficLight()
        self.connected_roads: List[Road] = []


class TrafficLight:
    def __init__(self):
        self.status = TrafficLightStatus.RED
        # 使用参数系统中的交通灯参数
        self.cycle_time = PARAMS["交通灯"]["cycle_time"]  # 周期时间(秒)
        self.green_time = PARAMS["交通灯"]["green_time"]  # 绿灯时间
        self.yellow_time = PARAMS["交通灯"]["yellow_time"]  # 黄灯时间

    def adjust_timing(self, road_densities: List[float]):
        """根据道路密度调整信号灯时间"""
        max_density = max(road_densities)
        self.green_time = min(45, max(15, int(30 * max_density)))
        self.cycle_time = self.green_time * 2

    def update(self, elapsed_time: float):
        """更新信号灯状态"""
        cycle_position = elapsed_time % self.cycle_time
        if cycle_position < self.green_time:
            self.status = TrafficLightStatus.GREEN
        elif cycle_position < self.green_time + self.yellow_time:  # 使用参数中的黄灯时间
            self.status = TrafficLightStatus.YELLOW
        else:
            self.status = TrafficLightStatus.RED


class TrafficSystem:
    def __init__(self):
        self.intersections: Dict[Position, Intersection] = {}
        self.roads: List[Road] = []
        self.vehicles: List[Vehicle] = []
        # 使用参数系统中的交通系统参数
        self.move_cooldown = PARAMS["交通系统"]["move_cooldown"]  # 移动冷却时间（秒）
        self.vehicle_spawn_rate = PARAMS["交通系统"]["vehicle_spawn_rate"]  # 车辆生成率
        self.max_vehicles = PARAMS["交通系统"]["max_vehicles"]  # 最大车辆数
        self.last_move_time = 0  # 上次移动时间
        self.visualization = None  # 添加可视化对象引用

    def add_intersection(self, position: Position):
        self.intersections[position] = Intersection(position)

    def add_road(self, start: Position, end: Position):
        road = Road(start, end)
        self.roads.append(road)
        if start in self.intersections:
            self.intersections[start].connected_roads.append(road)
        if end in self.intersections:
            self.intersections[end].connected_roads.append(road)

    def add_vehicle(self, vehicle: Vehicle):
        # 设置车辆默认速度
        vehicle.speed = PARAMS["车辆"]["default_speed"]
        self.vehicles.append(vehicle)
        self._calculate_route(vehicle)

    def _calculate_route(self, vehicle: Vehicle):
        """使用考虑道路密度的A*算法计算最优路径"""
        def heuristic(pos: Position) -> float:
            return abs(pos.x - vehicle.destination.x) + abs(pos.y - vehicle.destination.y)

        def get_road_density(pos1: Position, pos2: Position) -> float:
            """获取两个位置之间道路的密度"""
            for road in self.roads:
                if (road.start == pos1 and road.end == pos2) or (road.start == pos2 and road.end == pos1):
                    return road.density
            return float('inf')  # 如果没有道路连接，返回无穷大

        def get_neighbors(pos: Position) -> List[tuple[Position, float]]:
            """返回邻居节点及其代价，只返回通过道路连接的位置"""
            neighbors = []
            
            # 获取所有与当前位置相连的道路
            connected_roads = self.get_connected_roads(pos)
            
            for road in connected_roads:
                next_pos = road.end if road.start == pos else road.start
                # 使用道路长度和密度作为移动代价
                cost = road.length + road.density * 5  # 结合长度和密度
                neighbors.append((next_pos, cost))
            
            return neighbors

        start = vehicle.current_position
        goal = vehicle.destination
        
        # 重置路径信息
        vehicle.current_route_index = 0
        
        frontier = [(0, start)]
        came_from = {start: None}
        cost_so_far = {start: 0}
        
        while frontier:
            current_cost, current = heapq.heappop(frontier)
            
            if current == goal:
                break
            
            for next_pos, move_cost in get_neighbors(current):
                new_cost = cost_so_far[current] + move_cost
                
                if next_pos not in cost_so_far or new_cost < cost_so_far[next_pos]:
                    cost_so_far[next_pos] = new_cost
                    priority = new_cost + heuristic(next_pos)
                    heapq.heappush(frontier, (priority, next_pos))
                    came_from[next_pos] = current
        
        # 重建路径
        if goal in came_from:
            vehicle.route = self._reconstruct_path(came_from, start, goal)
        else:
            print(f"无法找到从 ({start.x}, {start.y}) 到 ({goal.x}, {goal.y}) 的路径")
            vehicle.route = [start]  # 设置一个默认路径

    def _reconstruct_path(self, came_from: Dict[Position, Position],
                          start: Position, goal: Position) -> List[Position]:
        current = goal
        path = []
        while current != start:
            path.append(current)
            current = came_from[current]
        path.append(start)
        path.reverse()
        return path

    def get_connected_roads(self, position: Position) -> List[Road]:
        """获取与给定位置相连的所有道路"""
        connected = []
        for road in self.roads:
            if road.start == position or road.end == position:
                connected.append(road)
        return connected

    def get_next_direction(self, vehicle: Vehicle) -> Direction:
        """根据车辆当前位置和路线确定下一个转向方向"""
        if not vehicle.route or len(vehicle.route) < 2:
            return Direction.STRAIGHT

        current_idx = vehicle.route.index(vehicle.current_position)
        if current_idx >= len(vehicle.route) - 1:
            return Direction.STRAIGHT

        current = vehicle.route[current_idx]
        next_pos = vehicle.route[current_idx + 1]

        # 简化版方向判断逻辑
        dx = next_pos.x - current.x
        dy = next_pos.y - current.y

        if dx > 0 and dy == 0:
            return Direction.RIGHT
        elif dx < 0 and dy == 0:
            return Direction.LEFT
        else:
            return Direction.STRAIGHT

    def update(self):
        """更新系统状态"""
        # 更新交通信号灯
        for intersection in self.intersections.values():
            connected_densities = [road.density for road in intersection.connected_roads]
            intersection.traffic_light.adjust_timing(connected_densities)

        # 更新车辆路径，但只在车辆不在移动时更新
        for vehicle in self.vehicles:
            if not vehicle.moving and random.random() < 0.1:  # 10%概率重新计算路径
                self._calculate_route(vehicle)

    def update_traffic_lights(self, elapsed_time: float):
        """更新所有交通信号灯状态"""
        for intersection in self.intersections.values():
            intersection.traffic_light.update(elapsed_time)

    def simulate_traffic(self):
        """模拟交通流量变化"""
        for road in self.roads:
            # 模拟车辆随机进出道路
            if random.random() < 0.1:  # 10%概率车辆数量变化
                change = random.randint(-2, 2)
                new_count = len(road.vehicles) + change
                road.vehicles = [Vehicle(f"sim_{i}", road.start, road.end) 
                               for i in range(max(0, min(new_count, road.capacity)))]

    def get_road_density(self, pos1: Position, pos2: Position) -> float:
        """获取两个位置之间道路的密度"""
        for road in self.roads:
            if (road.start == pos1 and road.end == pos2) or (road.start == pos2 and road.end == pos1):
                return road.density
        return float('inf')  # 如果没有道路连接，返回无穷大

    def has_road_connection(self, pos1: Position, pos2: Position) -> bool:
        """检查两个位置之间是否有道路连接"""
        return self.get_road_density(pos1, pos2) != float('inf')

    def move_vehicles(self, current_time: float):
        """移动所有车辆"""
        if current_time - self.last_move_time < self.move_cooldown:
            return
        
        self.last_move_time = current_time
        
        for vehicle in self.vehicles:
            if not vehicle.route or vehicle.current_route_index >= len(vehicle.route) - 1:
                if vehicle.current_position != vehicle.destination:
                    if vehicle.moving:  # 如果正在移动，等待到达当前目标点
                        continue
                    self._calculate_route(vehicle)
                continue
            
            # 如果车辆没有在移动，开始新的移动
            if not vehicle.moving:
                current_pos = vehicle.current_position
                next_pos = vehicle.route[vehicle.current_route_index + 1]
                
                # 验证是否存在道路连接
                if not self.has_road_connection(current_pos, next_pos):
                    if vehicle.moving:  # 如果正在移动，等待到达当前目标点
                        continue
                    self._calculate_route(vehicle)
                    continue
                
                # 检查交通灯
                if current_pos in self.intersections:
                    traffic_light = self.intersections[current_pos].traffic_light
                    if traffic_light.status == TrafficLightStatus.RED:
                        continue
                
                # 开始移动
                vehicle.moving = True
                vehicle.screen_x = current_pos.x
                vehicle.screen_y = current_pos.y
            
            # 如果车辆正在移动，更新位置
            if vehicle.moving:
                next_pos = vehicle.route[vehicle.current_route_index + 1]
                
                # 确保移动沿着道路进行
                dx = next_pos.x - vehicle.current_position.x
                dy = next_pos.y - vehicle.current_position.y
                
                # 计算单位向量
                length = (dx * dx + dy * dy) ** 0.5
                if length > 0:
                    dx = dx / length
                    dy = dy / length
                
                # 更新屏幕坐标
                move_speed = 0.1
                new_screen_x = vehicle.screen_x + dx * move_speed
                new_screen_y = vehicle.screen_y + dy * move_speed
                
                # 检查是否超出道路范围
                if self._is_position_on_road(vehicle.current_position, next_pos, new_screen_x, new_screen_y):
                    vehicle.screen_x = new_screen_x
                    vehicle.screen_y = new_screen_y
                
                # 检查是否到达目标位置
                if (abs(vehicle.screen_x - next_pos.x) < 0.1 and 
                    abs(vehicle.screen_y - next_pos.y) < 0.1):
                    vehicle.current_position = next_pos
                    vehicle.current_route_index += 1
                    vehicle.moving = False
                    
                    if vehicle.current_position == vehicle.destination:
                        if vehicle.id == "user_car":  # 只对用户车辆显示提示
                            self.visualization.show_arrival_notification()  # 触发到达提示
                        print(f"\n车辆已到达目的地：({vehicle.destination.x}, {vehicle.destination.y})")

    def _is_position_on_road(self, start: Position, end: Position, x: float, y: float) -> bool:
        """检查给定的坐标点是否在道路上"""
        # 计算道路的方向向量
        road_dx = end.x - start.x
        road_dy = end.y - start.y
        
        # 如果是垂直道路
        if road_dx == 0:
            return (abs(x - start.x) < 0.1 and 
                    min(start.y, end.y) <= y <= max(start.y, end.y))
        
        # 如果是水平道路
        if road_dy == 0:
            return (abs(y - start.y) < 0.1 and 
                    min(start.x, end.x) <= x <= max(start.x, end.x))
        
        return False

    def _get_random_destination(self, current_pos: Position) -> Position:
        """随机选择一个不同于当前位置的目的地"""
        positions = list(self.intersections.keys())
        positions.remove(current_pos)
        return random.choice(positions)

    def set_visualization(self, vis: TrafficVisualization):
        """设置可视化对象"""
        self.visualization = vis


# 使用示例
def main():
    # 初始化系统
    system = TrafficSystem()
    
    # 初始化可视化
    vis = TrafficVisualization(height=700)
    system.visualization = vis
    
    # 创建交叉路口网格
    for i in range(5):
        for j in range(5):
            system.add_intersection(Position(i, j))
    
    # 创建道路网络
    for i in range(5):
        for j in range(4):
            # 水平双向道路
            system.add_road(Position(i, j), Position(i, j + 1))
            system.add_road(Position(i, j + 1), Position(i, j))
            # 垂直双向道路
            system.add_road(Position(j, i), Position(j + 1, i))
            system.add_road(Position(j + 1, i), Position(j, i))
    
    while True:
        # 检查登录状态，未登录则显示登录界面
        if not vis.user_system.is_logged_in():
            login_result = vis.show_login_register_screen()
            if login_result == "quit" or login_result is False:
                vis.quit()
                return
            continue  # 登录成功后，继续显示主菜单
    
        # 已登录，显示主菜单
        menu_choice = vis.show_main_menu()
        
        # 判断是否需要注销
        if menu_choice == "logout":
            continue  # 注销后，返回登录界面
        
        # 处理菜单选择
        if menu_choice == "params":
            # 打开参数设置界面
            from parameters import open_params_window
            global PARAMS
            result = open_params_window()
            
            # 处理参数设置返回值
            if result == "back_to_main":
                # 如果是返回主菜单，则继续循环
                # 无需重新初始化pygame
                continue
            else:
                # 更新参数
                PARAMS = result
                # 无需重新初始化pygame
                continue
        elif menu_choice == "quit":
            vis.quit()
            return
        elif menu_choice != "start":
            continue  # 未知选项，重新显示主菜单
        
        # 通过UI获取起始点和目标点
        result = vis.get_input_positions()
        
        # 判断返回值类型，正确处理特殊返回值
        if isinstance(result, tuple) and len(result) == 2:
            # 正常的起点终点返回值
            start_pos, dest_pos = result
            if start_pos is None or dest_pos is None:
                vis.quit()
                return
        elif result == "back_to_main":
            # 返回主菜单
            continue
        else:
            # 其他情况，如直接退出
            vis.quit()
            return
        
        # 创建用户指定的车辆
        user_vehicle = Vehicle(
            id="user_car",
            current_position=start_pos,
            destination=dest_pos,
            speed=PARAMS["车辆"]["default_speed"]  # 使用参数中的默认速度
        )
        system.add_vehicle(user_vehicle)
        
        # 添加背景交通流量
        for road in system.roads:
            vehicle_count = random.randint(0, road.capacity // 2)
            road.vehicles = [
                Vehicle(f"background_{i}", road.start, road.end, speed=PARAMS["车辆"]["default_speed"])
                for i in range(vehicle_count)
            ]
        
        # 主循环
        start_time = time.time()
        try:
            while True:
                # 处理事件，检查是否需要返回主菜单
                event_result = vis.handle_events()
                if event_result == "back_to_main":
                    # 返回主菜单，重新启动主程序
                    break
                elif event_result is False:
                    # 用户关闭窗口，退出程序
                    vis.quit()
                    return
                    
                current_time = time.time() - start_time
                
                # 更新系统状态
                system.update()
                system.update_traffic_lights(current_time)
                system.simulate_traffic()
                system.move_vehicles(current_time)
                
                # 随机生成新车辆
                if random.random() < PARAMS["交通系统"]["vehicle_spawn_rate"] and len(system.vehicles) < PARAMS["交通系统"]["max_vehicles"]:
                    # 随机选择起点和终点
                    start_intersection = random.choice(list(system.intersections.keys()))
                    end_intersection = random.choice(list(system.intersections.keys()))
                    while end_intersection == start_intersection:
                        end_intersection = random.choice(list(system.intersections.keys()))
                    
                    # 创建新车辆
                    new_vehicle = Vehicle(
                        id=f"vehicle_{len(system.vehicles)}",
                        current_position=start_intersection,
                        destination=end_intersection,
                        speed=PARAMS["车辆"]["default_speed"]
                    )
                    system.add_vehicle(new_vehicle)
                
                # 渲染
                vis.render(system)
                
        except KeyboardInterrupt:
            pass


if __name__ == "__main__":
    main()

